from __future__ import annotations

import asyncio
import dataclasses
import time
from concurrent.futures import ThreadPoolExecutor

import ray

from ezflow import TraceStream as BaseStream
from ezflow.pipeline.base import BasePipeline


@dataclasses.dataclass
class RootConfig:
    name: str = "pipeline"
    version: str = "0.1"
    jdx: int = 0
    indices: list = dataclasses.field(default_factory=list)


class SimplePipeline(BasePipeline):
    @classmethod
    def from_config(cls, cfg: RootConfig) -> SimplePipeline:
        executor = ThreadPoolExecutor(max_workers=3)
        streams = []
        for i in cfg.indices:
            stream = SimpleStream(int, int, executor=executor)
            stream.idx = 3 * cfg.jdx + i
            streams.append(stream)
        instance: SimplePipeline = cls(*streams, with_lifetime=True)
        setattr(instance, "cfg", cfg)
        return instance


class SimpleStream(BaseStream):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.time = time.time()
        self.org_idx = 0

    def __call__(self, data):
        if data is None:
            return
        import random

        idx = getattr(self, "idx", 0) % 3
        time.sleep(idx)
        # 打印用时
        print(
            f"{getattr(self, 'idx', 0)}:{data} -> {time.localtime().tm_sec} <- {time.time() - self.time}"
        )
        self.time = time.time()
        return data

    async def async_send(self, nothing=None):
        await self.async_put(self.org_idx)
        self.org_idx += 1
        if self.org_idx > 10:
            await self.async_stop()


if __name__ == "__main__":
    import os

    os.environ["RAY_DEDUP_LOGS"] = "0"

    def create_pipeline(jdx=0):
        cfg = RootConfig(jdx=jdx, indices=[0, 1, 2])
        pipeline = SimplePipeline.from_config(cfg)
        return pipeline

    @ray.remote
    class PipelineActor:
        def __init__(self, jdx):
            self.pipeline = create_pipeline(jdx)  # 直接保存实例

        async def run(self):
            await self.pipeline.async_future()  # 调用实例的方法

    async def main():
        # 初始化 Ray
        ray.init(log_to_driver=True)

        # 创建两个 PipelineActor 实例
        actor1 = PipelineActor.remote(0)
        actor2 = PipelineActor.remote(1)

        # 异步运行两个 PipelineActor 的任务
        task1 = actor1.run.remote()
        task2 = actor2.run.remote()

        # 等待任务完成
        ray.wait([task1, task2])

        # 关闭 Ray
        ray.shutdown()

    asyncio.run(main())
